#!/usr/bin/env python
import sys, os, re

def showdescription():
	print """This script ckecks the load of several computers
Just feed it with a file of the type

computer1
computer2
...

and pass this file with the param -m <path to machinefile>.
(This file can be an mpi machinefile too). If no machinefile
was given the script looks in $HOME/machines for a file.

other params:
	-u :	gives detailed information on processes
	    	running on computers with load > 0.1
	-s : 	creates a machinefile of mpi-type 'freemachines'
	-? :	this text


(C) Stefan Fruhner 2008
"""


def gethstr(size):
	"""Get human readable string"""
	if size > 1000*1000*1000:
		return '%.2fGB' % (size/1024./1024./1024.)
	elif size > 1000*1000:
		return '%.2fMB' % (size/1024./1024.)
	elif size > 1000:
		return '%.2fkB' % (size/1024.)
	else:
		return '%.2fB' % (size)

	#if size > 1024*1024*1024:
		#return '%.2fGB' % (size/1024./1024./1024.)
	#elif size > 1024*1024:
		#return '%.2fMB' % (size/1024./1024.)
	#elif size > 1024:
		#return '%.2fkB' % (size/1024.)
	#else:
		#return '%.2fB' % (size)

class machine:
	def __init__(self, name):
		self.name		= name
		self.cpus		= 1
		self.top		= None
		self.load		= None
		self.short		= None
		self.overload	= False
		self.freecpus	= 1
		# values for free RAM
		self.MemTotal	= 0
		self.MemFree 	= 0
		self.Buffers	= 0
		self.Cached		= 0

	def __cmp__(self, other):
		return cmp(self.load, other.load)

#MYLOCAL = os.environ.get('MYLOCAL')
MYLOCAL = os.environ.get('HOME')
machinefile= MYLOCAL+'/machines'


arg	= sys.argv
l	= len(arg)

showhelp		= ('-?' in arg) #detailed view of user's processes
showusers		= ('-u' in arg) #detailed view of user's processes
savemachinefile	= ('-s' in arg) #save a machine file of mpi-type with free machines
#parsing params with value

if showhelp:
	showdescription()
	sys.exit(0)

if l > 1:
	for i in range(l):
		if arg[i]=='-m' and i+1 < l:
			machinefile = arg[i+1]

if not os.path.exists(machinefile):
	showdescription()
	sys.exit(0)

print "Using machinefile",machinefile

f = open(machinefile, 'r')
lines = f.readlines()
f.close()

#### get all machines and their processor numbers from the file
machines = []
longest_name = 0
for l in lines:
	if l.lstrip().startswith('#'): continue # ignore comments
	if l.find(':') > -1:
		try:
			name 		= l[:l.find(':')]
			processors 	= l[l.find(':')+1:].strip()
		except:pass
		else:
			if name: machines.append(machine(name)) # save
			if len(name) > longest_name: longest_name = len(name)
	else:
		n = l.strip()
		if n: machines.append(machine(n)) # save
		if len(n) > longest_name: longest_name = len(n)

longest_name+=1

print "visiting machines:"
for m in range(len(machines)):
	#print '\t',machines[m].name
	print '\r\t'+machines[m].name.ljust(longest_name,' '),
	sys.stdout.flush()
	sshscript = 'ssh %s "cat /proc/cpuinfo & cat /proc/meminfo & top -b -n 1 -d 0.1"' % (machines[m].name)
	res = os.popen(sshscript).readlines()
	for l in range(len(res)):
		res[l] = res[l].strip()
	machines[m].top = res

print '\nthinking about the load ... '

#getting average and short top view
for m in range(len(machines)):
	top = machines[m].top
	cnt = 0
	for l in top:
		# looking for the number of prcessors
		proc=re.findall("processor\s*:\s*([0-9]{1,2})$",l)
		if proc:
			machines[m].cpus = int(proc[0]) +1

		# looking for the free memory
		#MemFree:         16608 kB
		#Buffers:        290632 kB
		#Cached:        1142008 kB
		Mem=re.findall("MemTotal:\s*([0-9]*) kB$",l)
		if Mem:
			machines[m].MemTotal = int(Mem[0])

		Mem0=re.findall("^MemFree:\s*([0-9]*) kB$",l)
		if Mem0:
			machines[m].MemFree = int(Mem0[0])

		Mem1=re.findall("^Buffers:\s*([0-9]*) kB$",l)
		if Mem1:
			machines[m].Buffers = int(Mem1[0])

		Mem2=re.findall("^Cached:\s*([0-9]*) kB$",l)
		if Mem2:
			machines[m].Cached = int(Mem2[0])

		#load average: 0.03, 0.06, 0.02
		load=re.findall("load average: ([0-9\.]*), ([0-9\.]*), ([0-9\.]*)$",l)
		if load:
			machines[m].load = load

		#Mem:   4065808k total,  4043836k used,    21972k free,   255172k buffers
		mem=re.findall("Mem:\s*([0-9]*)k total,\s*([0-9]*)k used,\s*([0-9]*)k free.*",l)
		if mem:
			machines[m].memory = mem[0] # saves a list tot used free

		# looking for the first 5 lines of prcesses
		start = re.findall("^  PID",l)
		if re.match(r'.*PID.*', l):
			topfive = ''
			for i in range(5):
				topfive += machines[m].top[cnt+i]+'\n'
			machines[m].short = topfive
			#print machines[m].short
			break
		cnt += 1

for m in range(len(machines)):
	if not machines[m].load:
		continue
	load1   = float(machines[m].load[0][0])
	loadpercpu = load1/machines[m].cpus
	machines[m].load1 = load1
	machines[m].loadpercpu = loadpercpu
	free_assumed = int(float(machines[m].cpus) - load1)
	if free_assumed < 0 : free_assumed = 0;
	machines[m].freecpus = free_assumed;
	if float(machines[m].cpus) - load1 < 0.1:
		machines[m].overload = True

	if loadpercpu > 0.1 and showusers:
		print machines[m].name
		print 'load: %f (%f per cpu)' %(load1, loadpercpu)
		print machines[m].short,
		if machines[m].overload:
			print "Warning: maximum load is reached"
		print '\n'


machines.sort()
print 'machine'.ljust(20)+ ' | load  |CPUs|free|  tot RAM | '+'free'.rjust(10)
print 61*'-'

freecpus 	= 0
totcpus		= 0
totmem 		= 0
totfreemem 	= 0
totload		= 0.
for m in machines:
	if not m.load:
		machines.remove(m)
		continue
	m.FreeMem = m.MemFree + m.Buffers + m.Cached
	print '%s | %.3f | %s | %s | %s | %s' % (m.name.ljust(20), m.load1, str(m.cpus).rjust(2), str(m.freecpus).rjust(2), gethstr(int(m.MemTotal)*1024).rjust(8), gethstr(int(m.FreeMem)*1024).rjust(10) )
	totload		+= m.load1
	freecpus	+= m.freecpus
	totcpus		+= m.cpus
	totmem		+= m.MemTotal*1024
	totfreemem	+= m.FreeMem*1024
print 61*'_'
print '%s | %.3f | %01d | %d | %01s | %s' % ('all'.ljust(20), totload, totcpus, freecpus, gethstr(totmem).rjust(8), gethstr(totfreemem).rjust(10) )
print 61*'_'
print 61*'_'

if savemachinefile:
#	print "Writing machinefile 'freemachines'"
	l = []
	for m in machines:
		if not m.load: continue
		if m.freecpus > 0:
			l.append(m.name+':'+str(m.freecpus)+'\n')


	print "Writing machinefile 'freemachines'"
	f = open('freemachines', 'w+')
	f.writelines(l)
	f.close()






